import torch
from torch import nn
from net import MyAlexNet
import numpy as np
from torch.optim import lr_scheduler
import os

from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

import matplotlib.pyplot as plt

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

ROOT_TRAIN = r'C:\Users\Administrator\Desktop\Code\AlexNet\data\train'
ROOT_TEST = r'C:\Users\Administrator\Desktop\Code\AlexNet\data\val'

# 将图像的像素值归一化到【-1， 1】之间
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # 图像归一化 # 定义了一个归一化的转换操作

# 在训练集上进行数据预处理
train_transform = transforms.Compose([
    transforms.Resize((224, 224)), # 1. 图像调整大小
    transforms.RandomVerticalFlip(), # 2. 数据增强
    transforms.ToTensor(), # 3. 数据转换为张量数据
    normalize]) # 4. 使用之前定义的归一化操作对数据集进行处理

# 在验证集上进行数据预处理
val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize])
print('1. 数据集预处理完毕')

# 创建了训练集和验证集的数据集对象
train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
val_dataset = ImageFolder(ROOT_TEST, transform=val_transform)

# DataLoader类创建训练集和验证集的数据加载器。DataLoader是一个用于批量加载数据的实用类，它可以自动对数据进行分批、打乱和并行加载
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=True)
print('2. 数据集加载完成')

print('\n=========正在测试GPU===========')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if torch.cuda.is_available():
    device_count = torch.cuda.device_count()
    print(f"可用的GPU数量: {device_count}")

    for i in range(device_count):
        device_name = torch.cuda.get_device_name(i)
        print(f"GPU {i}: {device_name}")
else:
    print("没有可用的GPU")
model = MyAlexNet().to(device) # 将模型传入到device中进行计算，充分利用GPU的并行计算能力。
print('3. 模型已创建完毕')

# 定义一个损失函数
loss_fn = nn.CrossEntropyLoss()
print('4. 交叉熵损失函数已创建完毕')


# 定义一个优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
print('5. SGD优化器已创建完毕')


# 学习率每隔10轮变为原来的0.5
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

# 定义训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    for batch, (x, y) in enumerate(dataloader):
        image, y = x.to(device), y.to(device)
        output = model(image)
        cur_loss = loss_fn(output, y)
        _, pred = torch.max(output, axis=1)
        cur_acc = torch.sum(y==pred) / output.shape[0]

        # 反向传播
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()
        loss += cur_loss.item()
        current += cur_acc.item()
        n = n+1

    train_loss = loss / n
    train_acc = current / n
    print('train_loss' + str(train_loss))
    print('train_acc' + str(train_acc))
    return train_loss, train_acc

# 定义一个验证函数
def val(dataloader, model, loss_fn):
    # 将模型转化为验证模型
    model.eval()
    loss, current, n = 0.0, 0.0, 0
    with torch.no_grad():
        for batch, (x, y) in enumerate(dataloader):
            image, y = x.to(device), y.to(device)
            output = model(image)
            cur_loss = loss_fn(output, y)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(y == pred) / output.shape[0]
            loss += cur_loss.item()
            current += cur_acc.item()
            n = n + 1

    val_loss = loss / n
    val_acc = current / n
    print('val_loss' + str(val_loss))
    print('val_acc' + str(val_acc))
    return val_loss, val_acc

# 定义画图函数，绘制损失值变化趋势图
def matplot_loss(train_loss, val_loss):
    plt.plot(train_loss, label='train_loss')
    plt.plot(val_loss, label='val_loss')
    plt.legend(loc='best')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.title("训练集和验证集loss值对比图")
    plt.show()
# 定义画图函数，绘制准确率变化趋势图
def matplot_acc(train_acc, val_acc):
    plt.plot(train_acc, label='train_acc')
    plt.plot(val_acc, label='val_acc')
    plt.legend(loc='best')
    plt.ylabel('acc')
    plt.xlabel('epoch')
    plt.title("训练集和验证集acc值对比图")
    plt.show()



# 开始训练
loss_train = []
acc_train = []
loss_val = []
acc_val = []

epoch = 20
min_acc = 0
print('6. 超参数设定完成')

print('\n===========================模型开始训练！=============================')
for t in range(epoch):
    lr_scheduler.step()
    print(f"当前训练周期：【epoch{t+1}，{epoch}】\n-----------")
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer)
    val_loss, val_acc = val(val_dataloader, model, loss_fn)

    loss_train.append(train_loss)
    acc_train.append(train_acc)
    loss_val.append(val_loss)
    acc_val.append(val_acc)

    # 保存最好的模型权重
    if val_acc >min_acc:
        folder = 'save_model'
        if not os.path.exists(folder):
            os.mkdir('save_model')
        min_acc = val_acc
        print(f"save best model, 第{t+1}轮")
        torch.save(model.state_dict(), 'save_model/best_model.pth')
    # 保存最后一轮的权重文件
    if t == epoch-1:
        torch.save(model.state_dict(), 'save_model/last_model.pth')
print('===============模型训练完成！===================')

matplot_loss(loss_train, loss_val)
matplot_acc(acc_train, acc_val)
print('===============模型结果绘制完毕！===================')
